package edu.utah.seq.analysis;

import java.io.*;
import java.util.*;

import edu.utah.seq.data.*;
import trans.tpmap.WindowMaker;
import util.gen.*;

/**Use to find the peaks within EnrichedRegions*/
public class PeakPicker {
	//fields
	private EnrichedRegion[] enrichedRegions;
	private File[] treatmentPointDirs;
	private File[] controlPointDirs;
	private HashMap<String, PointData> treatments;
	private HashMap<String, PointData> controls;
	private int windowSize = 50;
	

	//constructors
	/**Picks the best peak based on rescanning each enriched region using the defined window size.*/
	public PeakPicker (EnrichedRegion[] enrichedRegions, File[] treatmentPointDirs, File[] controlPointDirs, int windowSize){
		this.enrichedRegions = enrichedRegions;
		this.treatmentPointDirs = treatmentPointDirs;
		this.controlPointDirs = controlPointDirs;
		this.windowSize = windowSize;
		pickPeaks();
	}
	
	
	//methods

	/**Main method of class. This method will then re window scan each EnrichedRegion and set the subWindows and bestSubWindow.*/
	public void pickPeaks(){

		//merge controls via concatination
		if (controlPointDirs != null) {
			controls = PointData.mergePointData(PointData.fetchPointData(controlPointDirs,null, false), false, true);
		}

		//merge treatments via concatination
		treatments = PointData.mergePointData(PointData.fetchPointData(treatmentPointDirs,null, false), false, true);

		//any controls?
		if (controls !=null){
			//subsample to match
			subSample();
			scoreTreatmentsControls();
		}
		//no controls
		else scoreTreatments();
		
		//find and assign best peak for each EnrichedRegion
		findPeaks();

	}

	/**Finds the best peak within each loaded enriched region.  Call scoreTreatmentsControls() or scoreTreatments() first.*/
	public void findPeaks(){
		for (int x=0; x< enrichedRegions.length; x++){
			SmoothingWindow bestSubWindow = findHighestPeak(enrichedRegions[x].getSubWindows(), 0);
			enrichedRegions[x].setBestSubWindow(bestSubWindow);
		}
	}
	
	/**Find the window with the largest score.  If multiple windows have the same best score then the start of the first and end
	 * of the last are used to create a new window and are returned.*/
	public static SmoothingWindow findHighestPeak(SmoothingWindow[] sm, int scoreIndex){
		if (sm == null || sm.length ==0) return null;
		//find highest score
		float bestScore = sm[0].getScores()[scoreIndex];
		for (int i=1; i< sm.length; i++){
			float testScore = sm[i].getScores()[scoreIndex];
			if (testScore > bestScore) bestScore = testScore;
		}
		//find smallest start and biggest end
		int start = sm[sm.length-1].getStart();
		int end = 0;
		for (int i=0; i< sm.length; i++){
			float testScore = sm[i].getScores()[scoreIndex];
			if (testScore == bestScore) {
				//set start?
				if (sm[i].getStart() < start) start = sm[i].getStart();
				//set end?
				if (sm[i].getStop() > end) end = sm[i].getStop();
			}
		}
		return new SmoothingWindow(start, end, new float[]{bestScore});
	}

	/**Used for window scanning with both treatment and control data.*/
	public void scoreTreatmentsControls(){
		//make window maker
		WindowMaker wm = new WindowMaker(windowSize,1);

		//for each enriched region, these are sorted by chromosome
		PointData treatPD = null; 
		PointData ctrlPD = null;
		int[] treatPos = null;
		int[] ctrlPos = null;
		String currentChromosome = "";
		for (int x=0; x< enrichedRegions.length; x++){
			//check chromosome and load PointData
			String testChromosome = enrichedRegions[x].getChromosome();
			if (currentChromosome.equals(testChromosome) == false){
				currentChromosome = testChromosome;
				//fetch PointData
				treatPD = treatments.get(currentChromosome);
				if (treatPD == null) Misc.printErrAndExit("\nError! Cannot find treatment PointData for chromosome -> "+currentChromosome+"\n");
				treatPos = treatPD.getPositions();
				ctrlPD = controls.get(currentChromosome);
				if (ctrlPD == null) Misc.printErrAndExit("\nError! Cannot find control PointData for chromosome -> "+currentChromosome+"\n");
				ctrlPos = ctrlPD.getPositions();
			}
			//find indexes		
			int[] treatIndexes = treatPD.findIndexes(enrichedRegions[x].getStart(), enrichedRegions[x].getStop());
			int[] controlIndexes = ctrlPD.findIndexes(enrichedRegions[x].getStart(), enrichedRegions[x].getStop());

			//fetch positions
			ArrayList<Integer> positions = new ArrayList<Integer>();
			for (int i=treatIndexes[0]; i< treatIndexes[1]; i++) positions.add(new Integer(treatPos[i]));
			for (int i=controlIndexes[0]; i< controlIndexes[1]; i++) positions.add(new Integer(ctrlPos[i]));
			int[] pos = Misc.integerArrayListToIntArray(positions);
			Arrays.sort(pos);

			//create windows
			int[][] windows = wm.makeWindows(pos);
			//assign bp positions
			for (int i=0; i< windows.length; i++){
				windows[i][0] = pos[windows[i][0]];
				windows[i][1] = pos[windows[i][1]]+1;	//last base isn't included
			}

			//for each window
			ArrayList<SmoothingWindow> smAL = new ArrayList<SmoothingWindow>();
			for (int i=0; i< windows.length; i++){

				//fetch scores
				float tSum = treatPD.sumPositionBP(windows[i][0], windows[i][1]);
				//don't proceed if tSum ==0
				if (tSum == 0) continue;
				float cSum = ctrlPD.sumPositionBP(windows[i][0], windows[i][1]);

				//calculate score, save window
				float score = (tSum- cSum) / new Double(Math.pow(tSum+cSum, 0.5)).floatValue();
				smAL.add(new SmoothingWindow (windows[i][0], windows[i][1], new float[]{score}));
			}
			//convert and add to enriched region
			SmoothingWindow[] sm = new SmoothingWindow[smAL.size()];
			smAL.toArray(sm);
			enrichedRegions[x].setSubWindows(sm);
		}

	}
	
	/**Used for window scanning treatment only datasets*/
	public void scoreTreatments(){
		//make window maker
		WindowMaker wm = new WindowMaker(windowSize,1);

		String currentChromosome = "";
		PointData treatPD = null;
		int[] treatPos = null;
		//for each enriched region, these are sorted by chromosome
		for (int x=0; x< enrichedRegions.length; x++){

			//check chromosome and load PointData
			String testChromosome = enrichedRegions[x].getChromosome();
			if (currentChromosome.equals(testChromosome) == false){
				currentChromosome = testChromosome;
				//fetch PointData
				treatPD = treatments.get(currentChromosome);
				if (treatments == null) Misc.printErrAndExit("\nError! Cannot find treatment PointData for chromosome -> "+currentChromosome+"\n");
				treatPos = treatPD.getPositions();
			}

			//find indexes
			int[] treatIndexes = treatPD.findIndexes(enrichedRegions[x].getStart(), enrichedRegions[x].getStop());

			//fetch positions
			ArrayList<Integer> positions = new ArrayList<Integer>();
			
			for (int i=treatIndexes[0]; i< treatIndexes[1]; i++) positions.add(new Integer(treatPos[i]));
			int[] pos = Misc.integerArrayListToIntArray(positions);

			//create windows
			int[][] windows = wm.makeWindows(pos);
			//assign bp positions
			for (int i=0; i< windows.length; i++){
				windows[i][0] = pos[windows[i][0]];
				windows[i][1] = pos[windows[i][1]]+1;	//last base isn't included
			}

			//for each window
			ArrayList<SmoothingWindow> smAL = new ArrayList<SmoothingWindow>();
			for (int i=0; i< windows.length; i++){

				//fetch scores
				float score = treatPD.sumPositionBP(windows[i][0], windows[i][1]);
				//don't proceed if tSum ==0
				if (score == 0) continue;

				//calculate score, save window
				smAL.add(new SmoothingWindow (windows[i][0], windows[i][1], new float[]{score}));
			}

			//convert and add to enriched region
			SmoothingWindow[] sm = new SmoothingWindow[smAL.size()];
			smAL.toArray(sm);
			enrichedRegions[x].setSubWindows(sm);
		}

	}


	/**Randomly tosses observations to match treatment and (half) control.*/
	private void subSample(){
		//count total observations
		int numTreatObs = PointData.totalObservations(treatments);
		int numCtrlObs = PointData.totalObservations(controls);

		//just match number of treatments to number of controls

		if (numTreatObs > numCtrlObs) {
			PointData.subSample(treatments, numCtrlObs);
		}
		else if (numTreatObs < numCtrlObs) {
			PointData.subSample(controls, numTreatObs);
		}

	}
	
}
